// A simple example of using read_graphviz to load a GraphViz Dot textual
// graph into a BGL adjacency_list graph

// Author: Ronald Garcia

#include "sim_util.h" //Must be included before Boost stuff to avoid Boost::Integer() namespace clobbering.

#include <boost/graph/graphviz.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/topological_sort.hpp>
#include <string>
#include <sstream>
#include <cstdio>
#include <cstdlib>
#include <iostream>
#include <fstream>
#include <vector>
#include <algorithm>
#include <omp.h> //For high-resolution timers
#include <cassert>
#include <sstream>

#include "RandomLib/Random.hpp"
#include "RandomLib/NormalDistribution.hpp"
#include "RandomLib/RandomSelect.hpp"
#include "types.h"
#include "hmetis_if.h"

#include "sim.h"

RandomLib::Random *RAND;
RandomLib::NormalDistribution<float> *NORMDIST;
RandomLib::RandomSelect<float> *SELECTDIST;

using namespace boost;
using namespace std;

mapping_strategy strat = hmetis;
bool commsweep = false;
float commsweepmin;
float commsweepmax;
float commsweepstep;
int commsweepiters;
float commdefault = 4.0;
bool sigmasweep = false;
float sigmasweepmin;
float sigmasweepmax;
float sigmasweepstep;
bool procsweep = false;
float procsweepmin;
float procsweepmax;
float procsweepstep;
int procsweepiters;
bool dumpoccsigma = false;
float dumpoccsigmaval;
bool dumpocccomm = false;
float dumpocccommval;
int iters = 1000;
int nprocs = -1;
std::string inputfilename;
std::string outputfileprefix;
std::string montecarlofilename;
std::string sigmasweepfilename;
std::string commsweepfilename;
std::string procsweepfilename;
std::string outputgraphfilename;

// My property tag classes

int main(int argc, char *argv[]) {
  RAND = new RandomLib::Random();
  NORMDIST = new RandomLib::NormalDistribution<float>();
  SELECTDIST = new RandomLib::RandomSelect<float>();
  hmetis_if hif;

  for(int i = 1; i < argc; i++) {
    std::string s(argv[i]);
    if(s.compare("-h") == 0) { //Do Hmetis
      strat = hmetis;
      continue;
    }
    if(s.compare("-g") == 0) { //Do Greedy
      strat = greedy;
      continue;
    }
    if(s.compare("-d") == 0) { //Do Dynamic
      strat = dynamic;
      continue;
    }
    if(s.compare("-in") == 0) { //Input filename
      assert(i < (argc-1));
      inputfilename = std::string(argv[++i]);
      continue;
    }
    if(s.compare("-out") == 0) { //Output file prefix
      assert(i < (argc-1));
      outputfileprefix = std::string(argv[++i]);
      continue;
    }
    if(s.compare("-procsweep") == 0) { //Proc sweep
      //Syntax: -procsweep <MINPROCS> <MAXPROCS> <STEPPROCS> <ITERS>
      assert(i < (argc-4));
      procsweep = true;
      procsweepmin = atoi(argv[++i]);
      assert(procsweepmin >= 1);
      procsweepmax = atoi(argv[++i]);
      assert(procsweepmax >= 1);
      procsweepstep = atoi(argv[++i]);
      assert(procsweepstep > 0.0);
      assert(procsweepmax >= procsweepmin);
      procsweepiters = atoi(argv[++i]);
      assert(procsweepiters > 0);
      continue;
    }
    if(s.compare("-commsweep") == 0) { //Sweep communication penalty
      assert(i < (argc-4));
      commsweep = true;
      commsweepmin = atof(argv[++i]);
      commsweepmax = atof(argv[++i]);
      commsweepstep = atof(argv[++i]);
      assert(commsweepstep > 0.0);
      assert(commsweepmax >= commsweepmin);
      commsweepiters = atoi(argv[++i]);
      assert(commsweepiters > 0);
      continue;
    }
    if(s.compare("-comm") == 0) { //Default communication penalty
      assert(i < (argc-1));
      commdefault = atof(argv[++i]);
      continue;
    }
    if(s.compare("-sigmasweep") == 0) { //Sweep fixed sigma (no iters necessary b/c it is deterministic)
      assert(i < (argc-3));
      sigmasweep = true;
      sigmasweepmin = atof(argv[++i]);
      sigmasweepmax = atof(argv[++i]);
      sigmasweepstep = atof(argv[++i]);
      assert(sigmasweepstep > 0.0);
      assert(sigmasweepmax >= sigmasweepmin);
      continue;
    }
    if(s.compare("-iters") == 0) { //Number of iterations
      assert(i < (argc-1));
      iters = atoi(argv[++i]);
      continue;
    }
    if(s.compare("-dumpoccsigma") == 0) { //Dump an occupancy graph for a particular fixed sigma value
      assert(i < (argc-1));
      dumpoccsigma = true;
      dumpoccsigmaval = atof(argv[++i]);
      continue;
    }
    if(s.compare("-dumpocccomm") == 0) { //Dump an occupancy graph for a particular communication penalty
      assert(i < (argc-1));
      dumpocccomm = true;
      dumpocccommval = atof(argv[++i]);
      continue;
    }
    if(s.compare("-p") == 0) { //# Of processors
      assert(i < (argc-1));
      std::string t(argv[++i]);
      if(t.compare("inf") == 0) { //Unlimited processors
        nprocs = -1;
      }
      else {
        nprocs = atoi(argv[i]);
      }
      continue;
    }
    if(s.compare("-imbalance") == 0) { //Allowed hMetis imbalance
      assert(i < (argc-1));
      hif.cfg.unbalance_factor = atof(argv[++i]);
      continue;
    }
    if(s.compare("-hmetisruns") == 0) { //# of hMetis runs
      assert(i < (argc-1));
      hif.cfg.num_runs = atoi(argv[++i]);
      continue;
    }
    assert(0 && "Should never reach, unknown command-line arg!");
  }

  if(inputfilename.length() == 0) {
    assert(0 && "Must specify input filename");
  }
  if(outputfileprefix.length() == 0) {
    assert(0 && "Must specify output file prefix");
  }

  ostringstream mc;
  mc << outputfileprefix << "." << nprocs << ".montecarlo";
  montecarlofilename = mc.str();
  ostringstream sig;
  sig << outputfileprefix << "." << nprocs << ".sigma";
  sigmasweepfilename = sig.str();
  ostringstream com;
  com << outputfileprefix << "." << nprocs << ".comm";
  commsweepfilename = com.str();
  ostringstream gra;
  gra << outputfileprefix << "." << nprocs << ".graph";
  outputgraphfilename = gra.str();
  ostringstream pro;
  pro << outputfileprefix << "." << nprocs << ".proc";
  procsweepfilename = pro.str();

  // Construct an empty graph and prepare the dynamic_property_maps.
  graph_t graph(0);
  dynamic_properties dp;
  
  property_map<graph_t, vertex_name_t>::type name =
  get(vertex_name, graph);
  dp.property("node_id",name);

  property_map<graph_t, vertex_color_t>::type color =
  get(vertex_color, graph);
  dp.property("color",color);

  property_map<graph_t, vertex_delay_mean_t>::type vdm =
  get(vertex_delay_mean, graph);
  dp.property("v_delay_mean",vdm);
  
  property_map<graph_t, vertex_delay_sd_t>::type vdsd =
  get(vertex_delay_sd, graph);
  dp.property("v_delay_sd",vdsd);

  property_map<graph_t, vertex_delay_dist_string_t>::type vdds =
  get(vertex_delay_dist_string, graph);
  dp.property("v_delay_dist",vdds);

//   property_map<graph_t, vertex_delay_dist_delays_t>::type vddd =
//   get(vertex_delay_dist_delays, graph);
//   dp.property("v_delay_dist_delays",vddd);
// 
//   property_map<graph_t, vertex_delay_dist_weights_t>::type vddw =
//   get(vertex_delay_dist_weights, graph);
//   dp.property("v_delay_dist_weights",vddw);
  
  property_map<graph_t, vertex_depth_vertices_t>::type vdepthv =
  get(vertex_depth_vertices, graph);
  dp.property("depth_v",vdepthv);

  property_map<graph_t, vertex_depth_time_t>::type vdeptht =
  get(vertex_depth_time, graph);
  dp.property("depth_t",vdeptht);
  
  property_map<graph_t, vertex_proc_t>::type vp =
  get(vertex_proc, graph);
  dp.property("proc",vp);
  
  property_map<graph_t, edge_delay_mean_t>::type edm =
  get(edge_delay_mean, graph);
  dp.property("e_delay_mean",edm);

  property_map<graph_t, edge_delay_sd_t>::type edsd =
  get(edge_delay_sd, graph);
  dp.property("e_delay_sd",edsd);

  property_map<graph_t, edge_delay_dist_string_t>::type edds =
  get(edge_delay_dist_string, graph);
  dp.property("e_delay_dist",edds);

//   property_map<graph_t, edge_delay_dist_delays_t>::type eddd =
//   get(edge_delay_dist_delays, graph);
//   dp.property("e_delay_dist_delays",eddd);
// 
//   property_map<graph_t, edge_delay_dist_weights_t>::type eddw =
//   get(edge_delay_dist_weights, graph);
//   dp.property("e_delay_dist_weights",eddw);

  // Use ref_property_map to turn a graph property into a property map
  boost::ref_property_map<graph_t*,std::string>
  gname(get_property(graph,graph_name));
  dp.property(std::string("name"),gname);

  std::ifstream graphfile(inputfilename.c_str(), std::ifstream::in);
  std::ofstream outfile(outputgraphfilename.c_str(), std::ofstream::out);

  FILE *tempe = stderr;
  freopen("/dev/null", "w",stderr);
  
  bool status = read_graphviz(graphfile,graph,dp,"node_id");

  //If we passed -p inf or -p -1, we now set nprocs to the number of vertices in the graph.
  //NOTE that in this case, we don't need to use hMetis, as the partitioning is trivial.
  if(nprocs <= 0) {
    nprocs = num_vertices(graph);
    graph_traits<graph_t>::vertex_iterator vi, vi_end;
    int i = 0;
    for(tie(vi, vi_end) = vertices(graph); vi != vi_end; ++vi) {
      put(vertex_proc, graph, *vi, i++);
    }
  }
  
  stderr = tempe;
  
  assert(status && "Did not read graph correctly");
  printf("Read graph successfully: %d vertices, %d edges\n", num_vertices(graph), num_edges(graph));
  sim_util::transformSelectDistributions(graph);

  if((strat == hmetis) && ((unsigned)nprocs < num_vertices(graph))) {
    hif.partition(graph, nprocs, outputfileprefix);
  }
  
  std::vector<Vertex> c;
  topological_sort(graph, std::back_inserter(c));

  sim::calc_depths(graph, c);

  std::vector<Vertex> inputs;
  std::vector<Vertex> outputs;

  graph_traits<graph_t>::vertex_iterator vi, vi_end;
  for(tie(vi, vi_end) = vertices(graph); vi != vi_end; ++vi) {
    if(in_degree(*vi, graph) == 0) //Input
      inputs.push_back(*vi);
    if(out_degree(*vi, graph) == 0) //Output
      outputs.push_back(*vi);
  }
  
  
  //The node indices are now in c in reverse topological order.
  //For now, each node has a vector of receive times from its predecessors.
  std::vector<double> times;
  times.resize(num_vertices(graph));

  //Number of iterations in which each vertex was on the critical path.
  std::vector<int> cpFrequency;
  cpFrequency.resize(num_vertices(graph));
  
  std::ofstream timing(montecarlofilename.c_str(), std::ofstream::out);


  double totaltime = 0.0f;
  double totalwork = 0.0f;

  double start, end;
  start = omp_get_wtime();
  
  //Monte carlo
  sim_mapping_strategy st;

  for(int i = 0; i < iters; i++)
  {
    if((iters > 10) && (i % (iters/10) == 0)) printf("Monte Carlo Simulation %d\n", i);
    switch(strat) {
      case hmetis: st = s_stored; break;
      case greedy: if(i == 0) st = s_greedy_store; else st = s_stored; break;
      case dynamic: st = s_greedy; break;
      default: assert(0); break;
    }
    sim_result sr = sim::simulate(graph, c, inputs, outputs, nprocs, (i == 0), false, 0.0f, commdefault, st);
    totaltime += sr.time;
    totalwork += sr.totalwork;
    for(std::list<Vertex>::const_iterator it = sr.criticalpath.begin(), end = sr.criticalpath.end();
         it != end; ++it) {
      cpFrequency[*it]++;
    }
    timing << sr.time << std::endl;
  }

  // Measuring the elapsed time
  end = omp_get_wtime();

  timing.close();
  
  // Time calculation (in seconds)
  printf("Monte Carlo: %f seconds per iteration\n", (end-start)/iters);
  printf("%f evaluations per second\n", (num_edges(graph)+num_vertices(graph))/((end-start)/iters));
  printf("Average Time: %f, Average Occupancy: %f\n", (totaltime/iters), (totalwork/nprocs)/(totaltime));

  int max_cp_freq = *(std::max_element(cpFrequency.begin(), cpFrequency.end()));
  
  for(tie(vi, vi_end) = vertices(graph); vi != vi_end; ++vi) {
    put(vertex_cp_freq, graph, *vi, (int)((float)cpFrequency[*vi]/max_cp_freq*8)+1);
  }
  
  cout << endl;

  
  if(sigmasweep) { //Sweep fixed sigma
    switch(strat) {
      case hmetis: st = s_stored; break;
      case greedy:
      case dynamic: st = s_greedy; break;
      default: assert(0); break;
    }
    std::ofstream sigma(sigmasweepfilename.c_str(), std::ofstream::out);
    start = omp_get_wtime();
    int ii = 0;
    for(float i = sigmasweepmin; i <= sigmasweep; i += 0.01f) {
      ii++;
      sim_result sr = sim::simulate(graph, c, inputs, outputs, nprocs, false, true, i, commdefault, st);
      printf("%f sigma: %f\n", i, sr.time);
      sigma << i << " " << sr.time << std::endl;
    }
    end = omp_get_wtime();
    printf("Sigma Sweep: %f seconds per iteration\n", (end-start)/ii);
    printf("%f evaluations per second\n", (num_edges(graph)+num_vertices(graph))/((end-start)/ii));
  }

  if(commsweep) {//Sweep communication penalty
    std::ofstream commpenalty(commsweepfilename.c_str(), std::ofstream::out);
    int ii = 0;
    start = omp_get_wtime();
    for(float i = commsweepmin; i <= commsweepmax; i += commsweepstep) {
      ii++;
      double sweeptime = 0.0f;
      double sweepwork = 0.0f;
      for(int j = 0; j < commsweepiters; j++) {
        switch(strat) {
          case hmetis: st = s_stored; break;
          case greedy: if(j == 0) st = s_greedy_store; else st = s_stored; break;
          case dynamic: st = s_greedy; break;
          default: assert(0); break;
        }
        sim_result sr = sim::simulate(graph, c, inputs, outputs, nprocs, false, false, 0.0f, i, st);
        sweeptime += sr.time;
        sweepwork += sr.totalwork;
      }
      printf("Comm. Penalty %f: Average Time %f, Average Occ. %f\n", i, sweeptime/commsweepiters, (sweepwork/nprocs)/sweeptime);
      commpenalty << i << " " << (sweeptime/commsweepiters) << std::endl;
    }
    end = omp_get_wtime();
    printf("Comm Sweep: %f seconds per iteration\n", (end-start)/(ii*commsweepiters));
    printf("%f evaluations per second\n", (num_edges(graph)+num_vertices(graph))/((end-start)/(ii*commsweepiters)));
  }


  if(procsweep) { //Sweep # of processors
    std::ofstream psw(procsweepfilename.c_str(), std::ofstream::out);
    start = omp_get_wtime();
    for(int i = procsweepmin; i <= procsweepmax; i += procsweepstep) {
      double sweeptime = 0.0f;
      double sweepwork = 0.0f;
      if((strat == hmetis) && ((unsigned)i < num_vertices(graph))) {
        hif.partition(graph, i, outputfileprefix);
      }
      for(int j = 0; j < procsweepiters; j++) {
        switch(strat) {
          case hmetis: st = s_stored; break;
          case greedy: if(j == 0) st = s_greedy_store; else st = s_stored; break;
          case dynamic: st = s_greedy; break;
          default: assert(0); break;
        }
        sim_result sr = sim::simulate(graph, c, inputs, outputs, i, false, false, 0.0f, commdefault, st);
        sweeptime += sr.time;
        sweepwork += sr.totalwork;
      }
      printf("%d Procs: Average Time %f, Average Occ. %f\n", i, (sweeptime/procsweepiters), (sweepwork/i)/sweeptime);
      psw << i << " " << (sweeptime/procsweepiters) << std::endl;
    }
    end = omp_get_wtime();
    psw.close();
  }

  //Write out the file
  
  //write_graphviz(outfile, graph, make_label_writer(name));
  //write_graphviz(outfile, graph, dp, std::string("node_id"));
  //write_graphviz(outfile, graph, delay_writer<property_map<graph_t, vertex_delay_mean_t>::type, property_map<graph_t, vertex_delay_sd_t>::type>(vdm, vdsd));
  //write_graphviz(outfile, graph, make_mean_delay_writer(vdm, vdsd));
  //write_graphviz(outfile, graph, make_delay_writer(vdm, vdsd, vp));
  write_graphviz(outfile, graph, my_vertex_writer<graph_t>(graph), my_edge_writer<graph_t>(graph));
  return 0;
}
